﻿/// Copyright (c) 2012 John Tear
///
/// Permission is hereby granted, free of charge, to any person obtaining a copy of this software 
/// and associated documentation files (the "Software"), to deal in the Software without restriction, 
/// including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, 
/// and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, 
/// subject to the following conditions:
///
/// The above copyright notice and this permission notice shall be included in all copies or substantial 
/// portions of the Software.
///
/// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT 
/// LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN 
/// NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 
/// WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 
/// SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
///

using System;
using System.Collections.Generic;
using System.Text;
using System.Timers;

using com.feelgoodsoftware.sharptftp.common;
using com.feelgoodsoftware.sharptftp.utilities;

namespace com.feelgoodsoftware.sharptftp
{
    /// <summary>
    /// Delegate for received client errors.
    /// </summary>
    /// <param name="Sender">The server throwing the event.</param>
    /// <param name="Reuqest">The request (including client ref).</param>
    /// <param name="ErrorMessage">The error message.</param>
    public delegate void OnClientErrorDelegate(Object Sender, TftpClientRequest Request, TftpErrorCode Code, String ErrorMessage);
    
    /// <summary>
    /// Delegate for client timeouts.
    /// </summary>
    /// <param name="Sender">The server throwing the event.</param>
    /// <param name="Request">The request (including client ref).</param>
    public delegate void OnClientTimeoutDelegate(Object Sender, TftpClientRequest Request);

    /// <summary>
    /// Delegate for files received from a client.
    /// </summary>
    /// <param name="Sender">The server throwing the event.</param>
    /// <param name="Request">The request (including client ref).</param>
    /// <param name="Filename">The name of the file to save to.</param>
    /// <param name="Data">The data received from the client.</param>
    public delegate void OnFileReceivedDelegate(Object Sender, TftpClientRequest Request, string Filename, Byte[] Data);

    /// <summary>
    /// Delegate for file requests from a client.
    /// </summary>
    /// <param name="Sender">The server throwing the event.</param>
    /// <param name="Request">The request (including client ref).</param>
    /// <param name="Filename">The filename being requested by the client.</param>
    public delegate void OnFileRequestDelegate(Object Sender, TftpClientRequest Request, string Filename);

    /// <summary>
    /// A class implementing a TFTP server - RFC 1350
    /// </summary>
    public class TftpServer
    {
        /// <summary>
        /// Thrown when an error packet is received from a tftp client.
        /// </summary>
        public event OnClientErrorDelegate OnClientError;
        /// <summary>
        /// Thrown when a tftp client has timed out.
        /// </summary>
        public event OnClientTimeoutDelegate OnClientTimeout;
        /// <summary>
        /// Thrown when a file has been recieved from a tftp client.
        /// </summary>
        public event OnFileReceivedDelegate OnFileReceived;
        /// <summary>
        /// Thrown when a tftp client has made a file request.
        /// </summary>
        public event OnFileRequestDelegate OnFileRequest;

        /// <summary>
        /// Dictionary to hold all client requests. The string key is the socket so it can be easily recognised during 
        /// the lock-step process of tftp.
        /// </summary>
        public Dictionary<string, TftpClientRequest> ClientRequests { get; set; }

        /// <summary>
        /// Gets or sets the amount of time (in miliseconds) that a tftp client can be silent before being regarded as timed out.
        /// The default is 2000 (2 seconds)
        /// </summary>
        public double ClientTimeout { get; set; }

        /// <summary>
        /// Specifies whether the server will negotiate the block size with tftp clients, default is true
        /// </summary>
        public bool AllowNegotiationOfBlockSize { get; set; }

        // the udp socket used for network stuff
        UDPListener _socket = new UDPListener();
        // we need a lock for the request dictionary (udp datagrams are being accepted asynchrounously)
        object _requestLock = new object();
        // a timer for checking that tftps clients haven't timed out
        Timer _timer = new Timer();

        /// <summary>
        /// Constructor
        /// </summary>
        /// <param name="port">The UDP port to start listening on</param>
        public TftpServer(int port)
        {
            Initialise(port);
        }

        /// <summary>
        /// Constructor, starts the Tftp server listening on UDP port 69
        /// </summary>
        public TftpServer()
        {
            Initialise(69);
        }

        // initialisation method for constructor(s)
        private void Initialise(int port)
        {
            #if DEBUG
            Console.WriteLine("TFTP Server Starting...");
#endif
            // start listening on udp port 69 (tftp)
            _socket.OnDataReceived += new DataReceivedEvent(_socket_OnDataReceived);
            _socket.StartListening(port);

            // create a new request dictionary
            ClientRequests = new Dictionary<string, TftpClientRequest>();

            // create a timer that checks for clients that have timed out (and remove them), say every second
            _timer = new Timer(1000);
            _timer.Elapsed += new ElapsedEventHandler(_timer_Elapsed);
            _timer.Enabled = true;
            _timer.Start();

            // set a default value for client timeout
            ClientTimeout = 2000;

            // set default for block size negotiation
            AllowNegotiationOfBlockSize = true;
        }

        void _timer_Elapsed(object sender, ElapsedEventArgs e)
        {
            if (ClientRequests.Count != 0)
            {
#if DEBUG
                Console.WriteLine("Clearing timed out requests...");
#endif
                // go through the dictionary and make sure the client requests haven't timed out
                ClearTimedOutRequests();
            }
        }

        // removes any requests that are deemed to have timed out
        void ClearTimedOutRequests()
        {
            bool fullCheck = true;

            // todo: it may be a good feature to recover from a timeout rather than just remove the request.
            // for example if a udp datagram is lost in transit then this may be the reason for timeout.
            lock (_requestLock)
            {
                foreach (TftpClientRequest request in ClientRequests.Values)
                {
                    if (request.LastPacketReceived.AddMilliseconds(ClientTimeout) < DateTime.Now)
                    {
                        // remove as it has timed out
#if DEBUG
                        Console.WriteLine("Removing request from: " + request.ToString() + " due to timeout: last packet(" + request.LastPacketReceived.ToString() + "), timeout(" + DateTime.Now.ToString() + "),  (" + request.LastPacketReceived.AddMilliseconds(ClientTimeout).ToString() + ")");
#endif

                        if (OnClientTimeout != null)
                        {
                            OnClientTimeout(this, request);
                        }

                        ClientRequests.Remove(request.ToString());
                        fullCheck = false;
                        break;
                    }
                }
            }

            // recursively check if we had to break the foreach
            if (fullCheck == false)
            {
                ClearTimedOutRequests();
            }
        }

        // called when data is recieved from a tftp client
        void _socket_OnDataReceived(object sender, DataEventArgs args)
        {
            // todo: log the incoming packet
#if DEBUG
            Console.WriteLine("Data received from " + args.RemoteEndPoint.Address + ": 0x" + ByteArray.ToHex(args.BytesReceived));
#endif

            // todo: here we need to ensure that the client is in the allowed access list
            // the access list will likely be a List<string> containing IP addresses and
            // will be passed in from the 'client software'

            // first we see if the client already has a pending request
            // the requests are unique by udp socket (ip:port)
            string socket = args.RemoteEndPoint.Address.ToString() + ":" + args.RemoteEndPoint.Port;

            // reference to the request, either new or from a previous packet from the client
            TftpClientRequest request;
            // make a packet
            TftpBasePacket packet = TftpBasePacket.MakePacket(args.BytesReceived);

            // if the key does not exist, add a new request
            if (!ClientRequests.ContainsKey(socket))
            {
                // first check to make sure that this is a NEW request, the first packet from a client should be
                // either a RRQ or WRQ (read or write request)
                if (packet.OpCode != TftpOpCode.RRQ && packet.OpCode != TftpOpCode.WRQ)
                {
                    // silently ignore and discard
#if DEBUG
                    Console.WriteLine("Incorrect OpCode in packet, discarding");
#endif
                    return;
                }

                // note: we are adding a request so the server can keep a status of the conversation between itself
                // and each client. This is due to tftp's lock-step process/nature. See RFC 1350 for details.
                request = new TftpClientRequest();
                request.Client = args.RemoteEndPoint;

                // lock the client request dictionary in case it is being purged (i.e. for timeouts)
                lock (_requestLock)
                {
#if DEBUG
                    Console.WriteLine("Creating a new client request for " + request.ToString());
#endif
                    // ClientRequest.ToString() returns the socket used in the key
                    ClientRequests.Add(request.ToString(), request);
                }
            }

            // ref the request
            request = ClientRequests[socket];

            // now we can process the packet
            // process based on the object type returned in the previous statement
            if (packet.GetType() == typeof(TftpReadWritePacket))
                if (packet.OpCode == TftpOpCode.RRQ)
                    ProcessReadPacket(packet, request);
                else
                    ProcessWritePacket(packet, request);
            else if (packet.GetType() == typeof(TftpDataPacket))
                ProcessDataPacket(packet, request);
            else if (packet.GetType() == typeof(TftpAckPacket))
                ProcessAckPacket(packet, request);
            else if (packet.GetType() == typeof(TftpErrorPacket))
                ProcessErrorPacket(packet, request);
        }

        // A read request is a request from the tftp client to read a file form the server.
        // the server should respond (if the file exists, and there are no errors) with a data packet 
        // with block number 1
        void ProcessReadPacket(TftpBasePacket packet, TftpClientRequest request)
        {
#if DEBUG
            Console.WriteLine("Processing RRQ packet for " + request.ToString());
#endif
            TftpReadWritePacket pkt = packet as TftpReadWritePacket;

            if (pkt != null)
            {
                // record that this is a read request (RRQ)
                request.RequestType = TftpOpCode.RRQ;
                request.Filename = pkt.Filename;

                // record that no block numbers have been received
                // and the packet time stamp has just been received
                request.LastBlockNumberReceived = 0;
                request.LastPacketReceived = DateTime.Now;

                // throw event to indicate to the 'software' that a read request has been made.
                if (OnFileRequest != null)
                {
                    OnFileRequest(this, request, pkt.Filename);
                }
            }
        }

        // A write request is a request from the tftp client to write a file to the server.
        // the first request should be responded to by an ack packet with block number 0
        void ProcessWritePacket(TftpBasePacket packet, TftpClientRequest request)
        {
#if DEBUG
            Console.WriteLine("Processing WRQ packet for " + request.ToString());
#endif
            TftpReadWritePacket pkt = packet as TftpReadWritePacket;

            if (pkt != null)
            {
                // record that this is a write request (RRQ)
                request.RequestType = TftpOpCode.RRQ;
                request.Filename = pkt.Filename;

                // record that no block numbers have been received
                // and the packet time stamp has just been received
                request.LastBlockNumberReceived = 0;
                request.LastPacketReceived = DateTime.Now;

                // send back an ack packet with block number 0
                TftpAckPacket returnPacket = new TftpAckPacket(0);
                _socket.SendDataGram(returnPacket.ToByteArray(), request.Client);

                // todo: log the packet that has just been sent
#if DEBUG
                Console.WriteLine("Sending packet (ACK) to " + request.ToString() + " : 0x" + ByteArray.ToHex(returnPacket.ToByteArray()));
#endif
            }
        }

        // a data packet is a client sending data following a write request, each data packet should be responded to
        // with an ack packet, acknowledging the block number that has been received.
        void ProcessDataPacket(TftpBasePacket packet, TftpClientRequest request)
        {
#if DEBUG
            Console.WriteLine("Processing DATA packet for " + request.ToString());
#endif
            TftpDataPacket pkt = packet as TftpDataPacket;

            if (pkt != null)
            {
#if DEBUG
                Console.WriteLine("Packet Block Number: " + pkt.BlockNumber + ", expecting : " + (int)(request.LastBlockNumberReceived + 1));
#endif
                // check that the block received is the correct block
                if (pkt.BlockNumber != request.LastBlockNumberReceived + 1)
                {
                    // return an error to the client
                    TftpErrorPacket errPacket = new TftpErrorPacket(TftpErrorCode.IllegalOperation, "Last packet sent in wrong order.");
                    _socket.SendDataGram(errPacket.ToByteArray(), request.Client);

#if DEBUG
                    Console.WriteLine("Sending packet (ERROR) to " + request.ToString() + " : 0x" + ByteArray.ToHex(errPacket.ToByteArray()));
#endif
                    return;
                }

                // we append the recieved data to the data field in the request
                request.DataBuffer = ByteArray.Combine(request.DataBuffer, pkt.Data);

                // acknowledge the packet
                TftpAckPacket returnPacket = new TftpAckPacket(pkt.BlockNumber);
                _socket.SendDataGram(returnPacket.ToByteArray(), request.Client);

#if DEBUG
                Console.WriteLine("Sending packet (ACK) to " + request.ToString() + " : 0x" + ByteArray.ToHex(returnPacket.ToByteArray()));
#endif

                // mark request up to date
                request.LastBlockNumberReceived = pkt.BlockNumber;
                request.LastPacketReceived = DateTime.Now;

                // if the data length less is than negotiated block size or zero then this is the last packet
                if (pkt.Data.Length < request.NegotiatedBlockSize)
                {
                    // throw event to tell 'software' that the file transfer is complete
                    if (OnFileReceived != null)
                    {
                        OnFileReceived(this, request, request.Filename, request.DataBuffer);
                    }

                    // mark as complete
                    request.RequestComplete = true;

                    // the request can be removed as it is complete
                    lock (_requestLock)
                    {
                        ClientRequests.Remove(request.ToString());
                    }
                }
            }
        }

        // an ack packet sent by a tftp client is acknowledgement of the data sent back 
        // from the server following a read request
        void ProcessAckPacket(TftpBasePacket packet, TftpClientRequest request)
        {
#if DEBUG
            Console.WriteLine("Processing ACK packet for " + request.ToString());
#endif
            TftpAckPacket pkt = packet as TftpAckPacket;

            if (pkt != null)
            {
                // the block number must be in order
                if (pkt.BlockNumber != request.LastBlockNumberReceived + 1)
                {
                    // return an error to the client
                    TftpErrorPacket errPacket = new TftpErrorPacket(TftpErrorCode.IllegalOperation, "Acknowledgement of packet received in wrong order.");
                    _socket.SendDataGram(errPacket.ToByteArray(), request.Client);

                    // todo: log the packet
#if DEBUG
                    Console.WriteLine("Sending packet (ERROR) to " + request.ToString() + " : 0x" + ByteArray.ToHex(errPacket.ToByteArray()));
#endif
                    return;
                }

                // record the block number received
                // and the packet time stamp has just been received
                request.LastBlockNumberReceived = pkt.BlockNumber;
                request.LastPacketReceived = DateTime.Now;

                // continue sending the data
                if (!request.RequestComplete)
                {
                    SendNextDataBlockToClient(request);
                }
                else
                {
                    // remove the request
                    lock (_requestLock)
                    {
#if DEBUG
                        Console.WriteLine("Request complete for " + request.ToString());
#endif
                        ClientRequests.Remove(request.ToString());
                    }
                }
            }
        }

        // an error packet is sent when there are issues, this is not likely to be sent from a tftp client
        void ProcessErrorPacket(TftpBasePacket packet, TftpClientRequest request)
        {
#if DEBUG
            Console.WriteLine("Processing ERROR packet for " + request.ToString());
#endif
            TftpErrorPacket pkt = packet as TftpErrorPacket;

            if (pkt != null)
            {
                // throw an event to the 'software' indicating an error on the client
                if (OnClientError != null)
                {
                    OnClientError(this, request, pkt.ErrorCode, pkt.ErrorMessage);
                }

                // mark as complete
                request.RequestComplete = true;

                // remove the request
                lock (_requestLock)
                {
                    ClientRequests.Remove(request.ToString());
                }
            }
        }

        private void SendNextDataBlockToClient(TftpClientRequest ClientRequest)
        {
            byte[] dataToSend = new byte[0];

            dataToSend = ClientRequest.GetNextDataBlockToSend();

            // note: the next block is +1
            TftpDataPacket returnPacket = new TftpDataPacket((short)(ClientRequest.LastBlockNumberReceived + 1), dataToSend);
            _socket.SendDataGram(returnPacket.ToByteArray(), ClientRequest.Client);

            // if we have sent a zero length data field we are complete
            if (dataToSend.Length == 0 || dataToSend.Length < ClientRequest.NegotiatedBlockSize)
            {
                ClientRequest.RequestComplete = true;
            }

            // todo: log the packet
#if DEBUG
            Console.WriteLine("Sending packet (DATA) to " + ClientRequest.ToString() + " : 0x" + ByteArray.ToHex(returnPacket.ToByteArray()));
#endif

        }

        /// <summary>
        /// Sends a file to the tftp client, this should be called following a tftp RRQ from a client.
        /// </summary>
        /// <param name="Data">The byte array of the file to send.</param>
        /// <param name="ClientRequest">The reference to the client requesting the data.</param>
        public void SendFileDataToClient(byte[] Data, TftpClientRequest ClientRequest)
        {
            // load the data into the client request data buffer
            ClientRequest.DataBuffer = Data;
            ClientRequest.LastBlockNumberReceived = 0;

            // start sending the data
            SendNextDataBlockToClient(ClientRequest);
        }

        /// <summary>
        /// Sends an error to the tftp client and also terminates the request.
        /// </summary>
        /// <param name="Code">The error code.</param>
        /// <param name="ErrorMessage">Error message to send to client.</param>
        /// <param name="ClientRequest">The client request.</param>
        public void SendErrorToClient(TftpErrorCode Code, string ErrorMessage, TftpClientRequest ClientRequest)
        {
            TftpErrorPacket packet = new TftpErrorPacket(Code, ErrorMessage);
            _socket.SendDataGram(packet.ToByteArray(), ClientRequest.Client);

#if DEBUG
            Console.WriteLine("Sending packet (ERROR) to " + ClientRequest.ToString() + " : 0x" + ByteArray.ToHex(packet.ToByteArray()));
#endif
        }
    }
}
